/**
  CNOK project, Anyang Normal University, IMP-CAS
  \class TASMatrix
  \brief This class calculate the eikonal S-matrix. It takes nucleus densities
  and nucleon-nucleon scattering amplitude (fNN) as input. The output is stored
  in arrays.   Since the S-matrix is fitted with sum of Gaussians, this class
  takes care of the fitting as well.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2020/07/08
  \date Last modified: 2020/10/06 by SUN Yazhou
  \copyright 2020-2023 SUN Yazhou
  \copyright CNOK project, Anyang Normal University, IMP-CAS
*/

#include <fstream>
#include <iostream>
#include <cstring>
#include <fstream>
#include "TAMath.h"
#include "TASMatrix.h"
#include "TAEikonalPhase.h"
#include "TAFit.h"
#include "TAInterpolate.h"

using std::ofstream;
using std::cout;
using std::endl;

#define sqr TAMath::sqr
#define polint TAInterpolateD::PolyInter

static const cdouble I(0., 1.); // imaginary number unit I

// ap, at: proj(targ) nucleon size factor: rho_n(r)=exp(-r^2/alphap^2)
TASMatrix::TASMatrix(int zP, int aP, int zT, int aT, double ek, const string &densP,
    const string &densT, double ap, double at) : fAlphajR(nullptr), fAlphajI(nullptr),
    fFitted(false), fSmxed(false){
  fNG = 40; fRL = 40.;
  fNB = 1200; fBMax = 120.; // NB: number of b bins; BMax is in fm
  fEikonalPhase = new TAEikonalPhase(zP, aP, zT, aT, ek, ap, at);
  fEikonalPhase->SetProjectileDensity(densP);
  fEikonalPhase->SetTargetDensity(densT);
  fB = new double[fNB]{}; fSmr = new double[fNB]{};
  fSmi = new double[fNB]{}; fSm2 = new double[fNB]{};
  SMatrix(fB, fSmr, fSmi, fSm2); // assign b, smr, smi and |sm|^2
} // end the constructor
// the ctor version using optical potentials
TASMatrix::TASMatrix(int zP, int aP, int zT, int aT, double ek, const string &opt,
    double ap, double at) : fAlphajR(nullptr), fAlphajI(nullptr), fFitted(false),
    fSmxed(false){
  fNG = 40; fRL = 40.;
  fNB = 1200; fBMax = 120.; // NB: number of b bins; BMax is in fm
  fEikonalPhase = new TAEikonalPhase(zP, aP, zT, aT, ek, ap, at);
  fEikonalPhase->SetOpticalPotential(opt);
  fB = new double[fNB]{}; fSmr = new double[fNB]{};
  fSmi = new double[fNB]{}; fSm2 = new double[fNB]{};
  SMatrix(fB, fSmr, fSmi, fSm2); // assign b, smr, smi and |sm|^2
} // end the constructor
TASMatrix::~TASMatrix(){
  if(fEikonalPhase) delete fEikonalPhase;
  if(fAlphajR) delete [] fAlphajR;
  if(fAlphajI) delete [] fAlphajI;
  if(fB) delete [] fB;
  if(fSmr) delete [] fSmr;
  if(fSmi) delete [] fSmi;
  if(fSm2) delete [] fSm2;
} // end the destructor

static const int NP = 6;
/// calculate and output the S-matrix
double TASMatrix::SMatrixR(double b) const{
  return polint(fB, fSmr, fNB, NP, b);
} // end member function SMatrixR
double TASMatrix::SMatrixI(double b) const{
  return polint(fB, fSmi, fNB, NP, b);
} // end member function SMatrixI
double TASMatrix::SMatrix2(double b) const{
  return polint(fB, fSm2, fNB, NP, b);
} // end member function SMatrixI
cdouble TASMatrix::SMatrix(double b) const{
  return SMatrixR(b) + I*SMatrixI(b);
} // end member function SMatrixI
///< real and imag part of S-matrix, in an array (0, kBMax) with step h=kBMax/fNB
void TASMatrix::SMatrix(double *b, double *smr, double *smi, double *sm2){
  if(fSmxed) return;
  fSmxed = true;

  char name[128] = "vis/sc.txt";
  if(1 == fEikonalPhase->fAP) sprintf(name, "vis/sv.txt");
  ofstream fout(name);
  cout << "Calulating S-Matrix to \033[32;1m" << name << "\033[0m..." << endl;
  const double h = fBMax/fNB, alphaNN = fEikonalPhase->GetAlphaNN();
  for(int i = 0; i < fNB; i++){ // chi(b) = u+iv //
    b[i] = i*h;
    // the (total) nuclear imaginary phase
    const double v = fEikonalPhase->GetPhaseNI(b[i]);
    // un: real nuclear eikonal phase; uc: Coulomb phase
    const double un = v*alphaNN, uc = fEikonalPhase->GetPhaseC(b[i]);
    const double u = un + uc; // the total real phase
    const double ev = exp(-v);
    smr[i] = ev*cos(u); smi[i] = ev*sin(u); sm2[i] = ev*ev;
    // cout << "b=" << b[i] << endl; // DEBUG
    // cout << "PSHIFT=(" << un << "," << v << ")" << endl; // DEBUG
    // cout << "Coulomb phase: " << uc << endl; // DEBUG
    // cout << "SCOULR=" << cos(uc) << ", SCOULI=" << sin(uc) << endl; // DEBUG
    // cout << "BEFORE Coulomb, " << " SM=("<< ev*cos(un) << "," << ev*sin(un) << ")\n"; // DEBUG
    // cout << "AFTER Coulomb, " << " SM=("<< smr[i] << "," << smi[i] << "),"; // DEBUG
    // cout << "|SM|^2=" << sm2[i] << endl; // DEBUG
    // getchar(); // DEBUG
    fout << b[i] << " " << smr[i] << " " << smi[i] << " " << sm2[i] << endl; // DEBUG
  } // end for over i
  fout.close();
  cout << "...DONE. S-Matrices can be inspected via";
  cout << " `cd vis/` and `root -l sm.C`" << endl;
  cout << "(provided ROOT is installed)." << endl << endl;
} // end member function SMatrix
/// fit the S-matrix to an expansion of Gaussians
/// the results are stored in member arrays alphaj (a complex number)
void TASMatrix::GaussianFit(){
  if(fFitted) return;
  fFitted = true;

  // calculate a list of S-Matrix w.r.t. b (impact parameter)
  // real and imag part of the s-matrix
  cout << "\033[32;1mFitting Sc with an expansion of Gaussians.";
  cout << " This may take a while...\033[0m" << endl;
  const int n = fNB, m = fNG; // nb: Nof bins of b; ng: Nof gaussians
  if(!fAlphajR && !fAlphajI){ fAlphajR = new double[m]; fAlphajI = new double[m]; }
  else TAException::Error("TASMatrix", "Gaussian: fAlphaj has been assigned.");

  bool isFit[m]; memset(isFit, 1, sizeof(isFit)); // true: for fit; false: fixed
  // alphaj for the 1-st gaussian is fixed to 1.
  // as usually S-Matrix approaches 1. for large b
  isFit[0] = false; fAlphajR[0] = 1.; fAlphajI[0] = 0.; // fix the first pair
  matrix covarR(m, m), covarI(m, m);
  double dy[n]; // dy: the variance of y (smr or smi); sm: S-Matrix
  double chi2R = 0., chi2I = 0., qR = 0., qI = 0.; // q: goodness-of-fit
  for(int i = 0; i < n; i++) dy[i] = 1.;
  TAFit::LSMFit(fB, fSmr, dy, n, m, fAlphajR, isFit, covarR, chi2R, gaus, &fRL, &qI);
  TAFit::LSMFit(fB, fSmi, dy, n, m, fAlphajI, isFit, covarI, chi2I, gaus, &fRL, &qR);
  // inspect the fitting quality //
  ofstream fout("vis/gausFit.txt");
  cout << "chi2R: " << chi2R << " chi2I: " << chi2I; // DEBUG
  cout << " qI: " << qI << " qR: " << qR << endl; // DEBUG
  for(int i = 0; i < fNB; i++){ // DEBUG
    double cr = 0., ci = 0.; // DEBUG
    for(int j = 0; j < fNG; j++){ // DEBUG
      cr += exp(-sqr(fB[i]*j/fRL))*fAlphajR[j]; // DEBUG
      ci += exp(-sqr(fB[i]*j/fRL))*fAlphajI[j]; // DEBUG
    } // DEBUG
    fout << fB[i] << " " << cr << " " << ci; // DEBUG
    fout << " " << fSmr[i] << " " << fSmi[i] << endl; // DEBUG
  } // DEBUG
  fout.close();
  cout << "\033[36;1mGauss Fit of Sc has been completed.";
  // cout << " The fitting can be inspected via" << endl;
  // cout << "       `cd vis/` and `root -l gausFit.C`\033[0m" << endl;
  // cout << "(provided ROOT is installed)." << endl;
} // end member function GaussianFit

double *TASMatrix::GetAlphajR(){
  if(!fAlphajR) GaussianFit();
  return fAlphajR;
} // end member function GetAlphajR
double *TASMatrix::GetAlphajI(){
  if(!fAlphajI) GaussianFit();
  return fAlphajI;
} // end member function GetAlphajI

// projectile mass in MeV
double TASMatrix::PM() const{ return fEikonalPhase->PM(); }
// target mass in MeV
double TASMatrix::TM() const{ return fEikonalPhase->TM(); }

/// gaussians used for S-Matrix expansion, \retval array funci[ma] evaluated at x
void TASMatrix::gaus(double x, double *funci, int ma, const double *p){
  for(int i = 0; i < ma; i++) funci[i] = exp(-sqr(x*i/p[0]));
} // end member function gaus
